import socket
import time
import random
import struct
import pandas as pd
import threading
from datetime import datetime

HEADER_FORMAT = 'BHHfB'#头部字节格式
HEADER_SIZE = struct.calcsize(HEADER_FORMAT)#头部字节大小

TYPE_NAMES = {
    1: "SYN", 2: "SYN-ACK", 3: "DATA", 4: "ACK", 5: "FIN"
}

class ReliableUDPClient:
    def __init__(self, server_ip, server_port, window_size_bytes=400, timeout=0.5):
        self.server_addr = (server_ip, server_port)
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)    #ipv4 + udp
        self.sock.settimeout(timeout) #设置套接字操作（如接收数据）的超时时间。

        # GBN协议参数
        self.window_size_bytes = window_size_bytes
        self.base = 0
        self.next_seq = 0
        self.packets = []   #存储待发送数据分片的列表

        # 网络状态跟踪
        self.estimated_rtt = 0.5  # 初始RTT估计（秒）,存储估计的往返时间。
        self.dev_rtt = 0.25  # RTT偏差
        self.timeout = timeout  # 动态超时
        self.rtt_list = []  #存储每次测量得到的 RTT 值

        # 统计信息
        self.send_times = {} #记录每个数据包的发送时间。字典的键是数据包的序列号（seq），值是该数据包发送时的时间戳
        self.retransmissions = 0
        self.sent_packets = 0
        self.total_packets = 0

        # 多线程控制
        self.lock = threading.Lock()    #确保在同一时间只有一个线程能够访问和修改共享资源
        self.stop_event = threading.Event() #用于控制 _send_thread 和 _recv_thread 线程的执行
        self.connection_established = False #记录客户端与服务器之间的连接状态

        # 快速重传相关
        self.dup_ack_count = 0  #记录连续收到的重复 ACK 的数量。
        self.last_ack = -1  #用于记录上一次收到的 ACK 号。

    def _update_timeout(self, sample_rtt_ms):
        """动态调整超时时间（RFC6298）"""
        sample_rtt = sample_rtt_ms / 1000  # 转换为秒
        self.estimated_rtt = 0.875 * self.estimated_rtt + 0.125 * sample_rtt #b = 0.25
        self.dev_rtt = 0.75 * self.dev_rtt + 0.25 * abs(sample_rtt - self.estimated_rtt) #a = 0.125
        self.timeout = max(0.1, self.estimated_rtt + 4 * self.dev_rtt) #确保超时时间不会小于 0.1 秒
        self.sock.settimeout(self.timeout)
        print(f"[超时调整] RTT={sample_rtt_ms:.2f}ms → 新超时={self.timeout:.3f}s")

    def connect(self):
        """三次握手建立连接"""
        print("\n===== 开始三次握手 =====")
        syn_packet = self._make_packet(1, 0, 0, b'') #构造一个 SYN 报文，

        for retry in range(5):
            self.sock.sendto(syn_packet, self.server_addr)
            print(f"握手尝试 {retry + 1}/5: 发送SYN")

            try:
                syn_ack, _ = self.sock.recvfrom(1024)
                if len(syn_ack) >= HEADER_SIZE: #检查接收到的数据长度是否大于等于报文头部的长度
                    ptype, _, ack, _, _ = struct.unpack(HEADER_FORMAT, syn_ack[:HEADER_SIZE]) #提取报文类型 ptype 和确认号 ack
                    if ptype == 2 and ack == 1:
                        print("收到SYN-ACK，发送ACK完成握手")
                        ack_packet = self._make_packet(4, 1, 1, b'')#构造一个 ACK 报文
                        self.sock.sendto(ack_packet, self.server_addr)
                        self.connection_established = True
                        return True
            except socket.timeout:
                print(f"握手超时，等待重试...")

        print("握手失败")
        return False

    def _send_thread(self):
        """发送线程实现"""
        # 新增：缓存上一次窗口状态的关键参数
        last_base = -1
        last_next = -1
        last_window_right = -1
        last_window_bytes = -1

        while not self.stop_event.is_set():
            with self.lock:
                if self.base >= self.total_packets:
                    break

                # 动态计算窗口边界（原有逻辑不变）
                window_bytes = 0
                window_right = self.base
                while (window_right < self.total_packets and
                       window_bytes + len(self.packets[window_right]) <= self.window_size_bytes):
                    window_bytes += len(self.packets[window_right])
                    window_right += 1

                # 核心优化：仅当窗口关键参数变化时才打印
                current_state_changed = (
                        self.base != last_base or
                        self.next_seq != last_next or
                        window_right != last_window_right or
                        window_bytes != last_window_bytes
                )

                if current_state_changed:
                    print(f"\n[窗口状态] base={self.base}, next={self.next_seq}, "
                          f"window=({self.base}-{window_right - 1}), 大小={window_bytes}B")
                    # 更新缓存的状态参数
                    last_base = self.base
                    last_next = self.next_seq
                    last_window_right = window_right
                    last_window_bytes = window_bytes

                # 发送窗口内未发送的包（原有逻辑不变）
                while self.next_seq < window_right:
                    packet = self._make_packet(3, self.next_seq, 0, self.packets[self.next_seq])
                    send_time = time.time()
                    self.sock.sendto(packet, self.server_addr)
                    self.send_times[self.next_seq] = send_time

                    start_byte = sum(len(p) for p in self.packets[:self.next_seq])
                    end_byte = start_byte + len(self.packets[self.next_seq])

                    print(f"[发送] seq={self.next_seq} bytes={start_byte}-{end_byte} "
                          f"size={len(self.packets[self.next_seq])}B time={send_time:.3f}")

                    self.next_seq += 1
                    self.sent_packets += 1

            time.sleep(0.1)  # 避免CPU忙等待

    def _recv_thread(self):
        """接收线程实现"""
        # 新增：记录上一次base，用于判断窗口是否滑动
        last_base = self.base
        while not self.stop_event.is_set():
            try:
                ack_packet, _ = self.sock.recvfrom(1024)
                if len(ack_packet) >= HEADER_SIZE:
                    ack_type, _, ack_num, _, _ = struct.unpack(HEADER_FORMAT, ack_packet[:HEADER_SIZE])

                    if ack_type == 4:  # ACK包
                        with self.lock:
                            # 优化1：仅在ACK_num与当前base不同时打印收到ACK（减少重复）
                            if ack_num != self.base or ack_num != self.last_ack:
                                print(f"\n[收到ACK] ack_num={ack_num}, 当前base={self.base}")

                            # 重复ACK检测
                            if ack_num == self.last_ack:
                                self.dup_ack_count += 1
                                # 优化2：重复ACK计数仅在达到阈值（3次）时打印，减少中间计数输出
                                if self.dup_ack_count >= 3:
                                    print(f"重复ACK计数: {self.dup_ack_count}")
                                    print("触发快速重传!")
                                    self._retransmit()
                            else:
                                self.last_ack = ack_num
                                self.dup_ack_count = 0

                            # 有效ACK处理
                            if ack_num >= self.base:
                                # RTT计算（仅使用按序ACK）
                                if ack_num == self.base and self.base in self.send_times:
                                    rtt = (time.time() - self.send_times[self.base]) * 1000
                                    self.rtt_list.append(rtt)
                                    self._update_timeout(rtt)
                                    print(f"[RTT测量] seq={self.base} RTT={rtt:.2f}ms")

                                # 滑动窗口
                                old_base = self.base
                                self.base = ack_num + 1
                                # 优化3：仅在base实际变化（窗口滑动）时打印窗口滑动信息
                                if self.base != last_base:
                                    print(f"[窗口滑动] base从{old_base}→{self.base}")
                                    last_base = self.base  # 更新记录的base

                                # 清理已确认包记录
                                for seq in list(self.send_times.keys()):
                                    if seq < self.base:
                                        del self.send_times[seq]

                                # 更新发送位置
                                self.next_seq = max(self.next_seq, self.base)

            except socket.timeout:
                with self.lock:
                    if self.base < self.total_packets:
                        print(f"\n[超时] base={self.base}, 最后RTT={self.rtt_list[-1] if self.rtt_list else 'N/A'}ms")
                        self._retransmit()

    def _retransmit(self):
        """执行重传"""
        self.retransmissions += 1
        print(f"重传开始: 从seq={self.base}重传窗口")

        # 重置next_seq并清除重复ACK计数
        self.next_seq = self.base
        self.dup_ack_count = 0
        self.last_ack = -1

        # 打印重传的字节范围
        start_byte = sum(len(p) for p in self.packets[:self.base])
        end_byte = sum(len(p) for p in self.packets[:self.base + 1])
        print(f"重传范围: bytes {start_byte}-{end_byte}")

    def send_data(self, data):
        """启动数据传输"""
        if not self.connection_established:
            print("连接未建立!")
            return

        print("\n===== 开始传输数据 =====")
        print(f"原始数据: {len(data)}字节")

        # 数据分片
        index = 0
        while index < len(data):
            chunk_size = random.randint(40, 80)
            self.packets.append(data[index:index + chunk_size])
            index += chunk_size

        self.total_packets = len(self.packets)
        print(f"分割为 {self.total_packets} 个数据包")
        print(f"平均包大小: {len(data) / self.total_packets:.1f}字节")
        print(f"初始窗口: {self.window_size_bytes}字节 (~{self.window_size_bytes / 60:.1f}个包)\n")

        # 启动线程
        send_thread = threading.Thread(target=self._send_thread)
        recv_thread = threading.Thread(target=self._recv_thread)

        send_thread.start()
        recv_thread.start()

        # 显示传输进度
        while self.base < self.total_packets and not self.stop_event.is_set():
            print(f"\r[进度] {self.base}/{self.total_packets} ({self.base / self.total_packets:.1%})", end="")
            time.sleep(0.5)

        # 清理
        self.stop_event.set()#停止线程
        send_thread.join()
        recv_thread.join()

        print("\n===== 传输完成 =====")
        self._print_stats()

    def _make_packet(self, type, seq, ack, data):
        """构造协议报文"""
        header = struct.pack(HEADER_FORMAT, type, seq, ack, time.time(), len(data))
        return header + data

    def _print_stats(self):
        """打印统计信息"""
        total_sent = self.sent_packets + self.retransmissions
        loss_rate = (self.retransmissions / total_sent) * 100 if total_sent > 0 else 0

        print("\n====== 统计信息 ======")
        print(f"总包数: {self.total_packets}")
        print(f"发送总数: {total_sent} (含{self.retransmissions}次重传)")
        print(f"丢包率: {loss_rate:.2f}%")

        if self.rtt_list:
            df = pd.DataFrame(self.rtt_list, columns=['RTT'])
            print("\nRTT统计:")
            print(f"平均值: {df['RTT'].mean():.2f}ms")
            print(f"最大值: {df['RTT'].max():.2f}ms")
            print(f"最小值: {df['RTT'].min():.2f}ms")
            print(f"标准差: {df['RTT'].std():.2f}ms")

        print("=====================")

    def close(self):
        if not self.connection_established:
            return

        print("\n===== 开始四次挥手 =====")

        # 第一次挥手：发送FIN
        fin_packet = self._make_packet(5, 0, 0, b'')
        for attempt in range(3):  # 重试机制
            self.sock.sendto(fin_packet, self.server_addr)
            print(f"尝试 {attempt + 1}/3: 发送FIN")

            try:
                # 第二次挥手：等待ACK
                ack, _ = self.sock.recvfrom(1024)
                if len(ack) >= HEADER_SIZE and struct.unpack('B', ack[:1])[0] == 4:
                    print("收到ACK，等待FIN...")

                    # 第三次挥手：等待FIN
                    fin, _ = self.sock.recvfrom(1024)
                    if len(fin) >= HEADER_SIZE and struct.unpack('B', fin[:1])[0] == 5:
                        print("收到FIN，发送最后ACK")

                        # 第四次挥手：发送ACK
                        self.sock.sendto(self._make_packet(4, 0, 0, b''), self.server_addr)
                        break

            except socket.timeout:
                print(f"尝试 {attempt + 1} 超时")

        self.sock.close()
        print("连接已关闭")


if __name__ == "__main__":
    server_ip = input("服务器IP: ")
    server_port = int(input("端口: "))
    client = ReliableUDPClient(server_ip, server_port)

    if client.connect():
        test_data = b'This is a test message for reliable UDP communication. ' * 10
        client.send_data(test_data)
        client.close()